package com.dson.gateway.filter;

import org.springframework.cloud.gateway.filter.GatewayFilter;
import org.springframework.cloud.gateway.filter.GatewayFilterChain;
import org.springframework.core.io.buffer.DataBuffer;
import org.springframework.http.HttpStatus;
import org.springframework.http.server.reactive.ServerHttpRequest;
import org.springframework.http.server.reactive.ServerHttpResponse;
import org.springframework.web.server.ServerWebExchange;
import reactor.core.publisher.Mono;

import java.net.InetSocketAddress;
import java.util.*;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.Executors;
import java.util.concurrent.TimeUnit;

public class RegisterGatewayFilter implements GatewayFilter {

    private Map<String, List<Long>> ipRequestTimesMap = new ConcurrentHashMap<>();

    private int rateMinutes;
    private int rateRequests;

    public RegisterGatewayFilter(int rateMinutes, int rateRequests) {
        this.rateMinutes = rateMinutes;
        this.rateRequests = rateRequests;

        Executors.newSingleThreadScheduledExecutor()
                .scheduleWithFixedDelay(() -> {
                    long currentTime = System.currentTimeMillis();
                    long startTime = currentTime - TimeUnit.MINUTES.toMillis(rateMinutes);
                    for (Map.Entry<String, List<Long>> entry : ipRequestTimesMap.entrySet()) {
                        entry.getValue().removeIf(t -> t < startTime);
                    }
                }, 60, 60, TimeUnit.SECONDS);
    }

    @Override
    public Mono<Void> filter(ServerWebExchange exchange, GatewayFilterChain chain) {
        ServerHttpRequest request = exchange.getRequest();
        ServerHttpResponse response = exchange.getResponse();

        String remoteIp = request.getHeaders().getFirst("X-Real-IP");
        if (remoteIp == null) {
            remoteIp = Optional.ofNullable(request.getRemoteAddress())
                    .map(InetSocketAddress::getHostString)
                    .orElse(null);
        }
        if (remoteIp == null) {
            return chain.filter(exchange);
        }
        List<Long> requestTimes = ipRequestTimesMap.get(remoteIp);
        if (requestTimes == null) {
            requestTimes = new ArrayList<>();
        }
        long currentTime = System.currentTimeMillis();
        long startTime = currentTime - TimeUnit.MINUTES.toMillis(rateMinutes);

        Iterator<Long> iterator = requestTimes.iterator();
        while (iterator.hasNext()) {
            long requestTime = iterator.next();
            if (requestTime < startTime) {
                iterator.remove();
            } else {
                break;
            }
        }
        if (requestTimes.size() > rateRequests) {
            response.setStatusCode(HttpStatus.SEE_OTHER);
            String data = "您频繁进⾏注册，请求已被拒绝";
            DataBuffer wrap = response.bufferFactory().wrap(data.getBytes());
            return response.writeWith(Mono.just(wrap));
        }
        requestTimes.add(currentTime);
        ipRequestTimesMap.put(remoteIp, requestTimes);
        return chain.filter(exchange);
    }
}
